import torch
import os
import sys

cur_path = os.path.abspath(os.path.dirname(__file__))
sys.path.append(os.path.join(cur_path, "..", "src"))

from doasr.models.configuration import DoAsrConfig
from doasr.models.doasrformer.llm import DoAsrForCausalLM
from doasr.models.tokenizer import DoAsrTokenizer


model_name_or_path = os.path.join(cur_path, "..", "models_pretraining", "DoASR-0.5B-pretraining_20240907101412")
model = DoAsrForCausalLM.from_pretrained(model_name_or_path)
model = model.eval()
tokenizer = DoAsrTokenizer.from_pretrained(model_name_or_path)

prompt = "<|endoftext|>Anarchism is a political philosophy and movement that is sceptical of"

inputs = tokenizer(prompt, return_tensors="pt")
inputs = model.prepare_inputs_for_generation(inputs.input_ids)
result = model(**inputs)
print(result)